use super::{ffi, CudaStream, Error, Handle, Image, ImageInfo, JpegState, PixelFormat, Result};
use std::convert::TryInto;

pub struct Decoder {
    state: JpegState,
    handle: Handle,
    cuda: CudaStream,
}

impl Decoder {
    pub fn new() -> Result<Self> {
        let handle = Handle::new()?;
        let state = JpegState::new(&handle)?;
        let cuda = CudaStream::new()?;

        Ok(Self {
            state,
            handle,
            cuda,
        })
    }

    pub fn decode(&self, data: &[u8], fmt: PixelFormat) -> Result<Image> {
        use ffi::nvjpegStatus_t::NVJPEG_STATUS_SUCCESS;

        let info = self.get_image_info(data)?;
        let image = Image::new(info, fmt)?;
        let mut dest = image.to_nvjpeg_image();
        let status = unsafe {
            ffi::nvjpegDecode(
                self.handle.as_ptr(),
                self.state.as_ptr(),
                data.as_ptr(),
                data.len().try_into().unwrap(),
                fmt.into(),
                &mut dest,
                self.cuda.as_ptr(),
            )
        };
        match status {
            NVJPEG_STATUS_SUCCESS => Ok(image),
            status => Err(Error::Jpeg(status)),
        }
    }

    pub fn get_image_info(&self, data: &[u8]) -> Result<ImageInfo> {
        use super::MAX_COMPONENT;
        use ffi::nvjpegChromaSubsampling_t::NVJPEG_CSS_UNKNOWN;
        use ffi::nvjpegStatus_t::NVJPEG_STATUS_SUCCESS;

        let mut components = 0;
        let mut subsampling = NVJPEG_CSS_UNKNOWN;
        let mut widths: [i32; MAX_COMPONENT] = [0, 0, 0, 0];
        let mut heights: [i32; MAX_COMPONENT] = [0, 0, 0, 0];

        let status = unsafe {
            ffi::nvjpegGetImageInfo(
                self.handle.as_ptr(),
                data.as_ptr(),
                data.len() as u64,
                &mut components,
                &mut subsampling,
                widths.as_mut_ptr(),
                heights.as_mut_ptr(),
            )
        };

        match status {
            NVJPEG_STATUS_SUCCESS => {
                let widths = [
                    widths[0] as usize,
                    widths[1] as usize,
                    widths[2] as usize,
                    widths[3] as usize,
                ];
                let heights = [
                    heights[0] as usize,
                    heights[1] as usize,
                    heights[2] as usize,
                    heights[3] as usize,
                ];
                Ok(ImageInfo {
                    components: components as usize,
                    subsampling: subsampling.into(),
                    widths,
                    heights,
                })
            }
            status => Err(Error::Jpeg(status)),
        }
    }
}
